import torch


class FeatureTransform:
    """Base callable for feature transforms that operate on continuous covariates.
    Input shape: [B, T, N, D_cont]
    """
    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        return x


class ScaleStd(FeatureTransform):
    """Standardize continuous features per-feature using batch-time-node statistics.
    """
    def __init__(self, eps: float = 1e-6):
        self.eps = eps

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(dim=(0, 1, 2), keepdim=True)
        std = x.std(dim=(0, 1, 2), keepdim=True)
        return (x - mean) / (std + self.eps)


class ClipOutliers(FeatureTransform):
    """Clip extreme values to a specified range based on percentile.
    """
    def __init__(self, p_low: float = 1.0, p_high: float = 99.0):
        self.p_low = p_low
        self.p_high = p_high

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        # Compute percentiles on CPU for robustness
        x_cpu = x.detach().cpu()
        low = torch.quantile(x_cpu, self.p_low / 100.0, dim=(0, 1, 2), keepdim=True)
        high = torch.quantile(x_cpu, self.p_high / 100.0, dim=(0, 1, 2), keepdim=True)
        low = low.to(x.device)
        high = high.to(x.device)
        return torch.clamp(x, min=float(low.min()), max=float(high.max()))


class SemDropout(FeatureTransform):
    """Apply dropout to continuous features only during training to encourage robustness.
    """
    def __init__(self, p: float = 0.1):
        self.p = p

    def __call__(self, x: torch.Tensor) -> torch.Tensor:
        # Use torch.nn.functional.dropout with training flag from outer scope if set later
        return torch.nn.functional.dropout(x, p=self.p, training=True)


def build_feature_transforms(transforms_str: str):
    """Parse a comma-separated transform spec into a list of callables.
    Supported names: 'covar_norm' (ScaleStd), 'sem_dropout' (SemDropout), 'clip_outliers'.
    """
    if not transforms_str:
        return []
    name_map = {
        'covar_norm': ScaleStd(),
        'sem_dropout': SemDropout(p=0.1),
        'clip_outliers': ClipOutliers(),
    }
    out = []
    for name in [s.strip() for s in transforms_str.split(',') if s.strip()]:
        if name in name_map:
            out.append(name_map[name])
    return out